package benchmark.energy.server;

import benchmark.energy.EnergyDataTxGenerator;
import benchmark.energy.tx.*;
import benchmark.server.Neo4jSocketServer;
import benchmark.transaction.definition.AbstractTransaction;
import benchmark.utils.*;
import ch.qos.logback.core.BasicStatusManager;
import com.alibaba.fastjson.JSON;
import com.google.common.base.Preconditions;
import org.neo4j.graphdb.*;

import java.io.File;
import java.io.IOException;
import java.util.*;

public class Neo4jServer2 extends Neo4jSocketServer.ReqExecutor {

    // road name -> road id.
    private HashMap<Integer, Long> name2Id = new HashMap<>();
    // rId -> neo4j road node id.
    private HashMap<Long, Long> rId2Id = new HashMap<>();
    // road id -> latest neo4j relation node id.
    private HashMap<Long, Long> roadLatestRelId = new HashMap<>();
    // crossId -> neo4j cross node id.
    private HashMap<Long, Long> crossId2Id = new HashMap<>();

    public static void main(String[] args) {
        Neo4jSocketServer server = new Neo4jSocketServer(dbDir(), new Neo4jServer2());
        RuntimeEnv env = RuntimeEnv.getCurrentEnv();
        String serverCodeVersion = env.name() + "." + Helper.codeGitVersion();
        System.out.println("server code version:" + serverCodeVersion);
        try {
            server.start();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }


    private static File dbDir() {
        String path = Helper.mustEnv("DB_PATH");
        Preconditions.checkNotNull(path, "need arg: DB_PATH");
        File dbDir = new File(path);
        if (!dbDir.exists()) {
            if (dbDir.mkdirs()) return dbDir;
            else throw new IllegalArgumentException("invalid dbDir");
        } else if (!dbDir.isDirectory()) {
            throw new IllegalArgumentException("invalid dbDir");
        }
        return dbDir;
    }

    private enum Edge implements RelationshipType {
        REACH_TO, PROP;
    }

    @Override
    protected void setDB(GraphDatabaseService db) {
        this.db = db;
//        try (Transaction tx = db.beginTx()) {
//            Schema schema = db.schema();
//            schema.indexFor(Label.label("Road")).on("name").create();
//            schema.indexFor(Label.label("Cross")).on("id").create();
//            tx.success();
//        }
    }

    @Override
    protected AbstractTransaction.Result execute(String line) throws RuntimeException {
        AbstractTransaction tx = JSON.parseObject(line, AbstractTransaction.class);
        switch (tx.getTxType()) {
            case TX_IMPORT_STATIC_DATA:
                return execute((ImportStaticDataTx) tx);
            case TX_IMPORT_TEMPORAL_DATA:
                return execute((ImportTemporalDataTx) tx);
            case TX_UPDATE_TEMPORAL_DATA:
                return execute((UpdateTemporalDataTx) tx);
            case TX_QUERY_SNAPSHOT:
                return execute((SnapshotQueryTx) tx);
            case TX_QUERY_SNAPSHOT_AGGR_MAX:
                return execute((SnapshotAggrMaxTx) tx);
            case TX_QUERY_SNAPSHOT_AGGR_DURATION:
                return execute((SnapshotAggrDurationTx) tx);
            case TX_ENTITY_TEMPORAL_CONDITION_QUERY:
                return execute((EntityTemporalConditionTx) tx);
            default:
                throw new UnsupportedOperationException();
        }
    }


    Label propLabel = Label.label("PROP");
    private void createPropNode(Node n) {
        for(String propName : EnergyDataTxGenerator.properties){
            Node prop = db.createNode(propLabel);
            prop.setProperty("name", propName);
            Relationship r = n.createRelationshipTo(prop, Edge.PROP);
            r.setProperty("name", propName);
        }
    }

    private AbstractTransaction.Result execute(ImportStaticDataTx tx) {
        Label crossLabel = Label.label("NODE");
        try (Transaction t = db.beginTx()) {
            Map<Integer, Long> idMap = new HashMap<>();
            for (ImportStaticDataTx.EnergyNode node : tx.getNodes()) {
                Node n = db.createNode(crossLabel);
                n.setProperty("id", node.getRawId());
                n.setProperty("name", node.getName());
                n.setProperty("country", node.getCountry());
                n.setProperty("voltage", node.getVoltage());
                n.setProperty("latitude", node.getLatitude());
                n.setProperty("longitude", node.getLongitude());
                idMap.put(node.getRawId(), n.getId());
//                createPropNode(n);
            }
            for (ImportStaticDataTx.EnergyRel rel : tx.getRels()) {
                Node s = db.getNodeById(idMap.get(rel.getFromRawId()));
                Node e = db.getNodeById(idMap.get(rel.getToRawId()));
                Relationship r = s.createRelationshipTo(e, Edge.REACH_TO);
                r.setProperty("x", rel.getX());
                r.setProperty("y", rel.getY());
                r.setProperty("limit", rel.getLimit());
                r.setProperty("length", rel.getLength());
            }
            t.success();
        }
        return new AbstractTransaction.Result();
    }

    private void setProperties(List<Float> values, Node r){
        for(int i=0; i< EnergyDataTxGenerator.properties.length; i++) {
            String propName = EnergyDataTxGenerator.properties[i];
            Float val = values.get(i);
            r.setProperty(propName, val);
        }
    }

    private Node createTemporalPropNode(int id, int startTime, int endTime) {
        Node ret;
        ResourceIterator<Node> nodes = db.findNodes(Label.label("TPNODE"), "node_id", id, "st_time", startTime, "en_time", endTime);
        if (nodes.hasNext()) {
            ret = nodes.next();
        } else {
            ret = db.createNode(Label.label("TPNODE"));
            ret.setProperty("node_id", id);
            ret.setProperty("st_time", startTime);
            ret.setProperty("en_time", endTime);
        }
        nodes.close();
        return ret;
    }

    private AbstractTransaction.Result execute(ImportTemporalDataTx tx) {
        try (Transaction t = db.beginTx()) {
            for (ImportTemporalDataTx.EnergyStatus s : tx.getData()) {
                Long id = name2Id.get(s.getNodeId());
                assert id != null;
                Node n = db.getNodeById(id);
                int time = s.getTimePoint();
                Node tp = createTemporalPropNode(s.getNodeId(), time, time+3599);
                setProperties(s.getStatus(), tp);
                n.createRelationshipTo(tp, Edge.PROP);
            }
            t.success();
        }
        return new AbstractTransaction.Result();
    }


    private void copyProperty(Node from, Node to){
        for(String propName : EnergyDataTxGenerator.properties) {
            to.setProperty(propName, from.getProperty(propName));
        }
    }

    private AbstractTransaction.Result execute(UpdateTemporalDataTx tx) {
        int tBegin = tx.getStartTime();
        int tFinish = tx.getEndTime();
        try (Transaction t = db.beginTx()) {
            // Do not consider that the old value equals to the new value
            Node node = db.findNode(Label.label("NODE"), "id", tx.getEntityId());
            HashSet<Node> willBeDeleted = new HashSet<>();
            for(Relationship r: node.getRelationships(Edge.PROP)){
                Node tp = r.getOtherNode(node);
                int st = (int) tp.getProperty("st_time");
                int en = (int) tp.getProperty("en_time");
                if (tFinish >= st && tBegin <= en) { // has overlap
                    if (st < tBegin && tFinish < en) { // should split node.
                        Node theOtherPart = db.createNode(Label.label("TPNODE"));
                        tp.setProperty("en_time", tBegin-1);
                        copyProperty(tp, theOtherPart);
                        theOtherPart.setProperty("st_time", tFinish+1);
                        theOtherPart.setProperty("en_time", en);
                        theOtherPart.setProperty("node_id", tx.getEntityId());
                        node.createRelationshipTo(theOtherPart, Edge.PROP);
                    }else if(tBegin <= st && en <= tFinish){ // should delete
                        willBeDeleted.add(tp);
                    }else if(tBegin <= st && tFinish < en ){ // change start
                        tp.setProperty("st_time", tx.getEndTime() + 1 );
                    }else if(st < tBegin && en <= tFinish){ // change end
                        tp.setProperty("en_time", tx.getStartTime() - 1 );
                    }
                }
            }
            for(Node toDel : willBeDeleted){
                for(Relationship r : toDel.getRelationships()){
                    r.delete();
                }
                toDel.delete();
            }
            Node insertedTp = db.createNode(Label.label("TPNODE"));
            insertedTp.setProperty("node_id", tx.getEntityId());
            insertedTp.setProperty("st_time", tx.getStartTime());
            insertedTp.setProperty("en_time", tx.getEndTime());
            setProperties(tx.getValues(), insertedTp);
            t.success();
        }
        return new AbstractTransaction.Result();
    }

    private AbstractTransaction.Result execute(SnapshotQueryTx tx) {
        try (Transaction t = db.beginTx()) {
            List<Pair<String, Float>> res = new ArrayList<>();
            int time = tx.getTimestamp();
            try(ResourceIterator<Node> it = db.findNodes(Label.label("TPNODE"))) {
                while(it.hasNext()) {
                    Node tp = it.next();
                    int startTime = (int) tp.getProperty("st_time");
                    int endTime = (int) tp.getProperty("en_time");
                    if (startTime <= time && time <= endTime) {
                        String name = null;
                        for(Relationship r : tp.getRelationships()){
                            Node main = r.getOtherNode(tp);
                            name = (String) main.getProperty("name");
                        }
                        assert name!=null;
                        float v = (Float) tp.getProperty(tx.getPropertyName());
                        res.add(Pair.of(name, v));
                    }
                }
            }
            t.success();
            SnapshotQueryTx.Result result = new SnapshotQueryTx.Result();
            result.setStatus(res);
            return result;
        }
    }

    private AbstractTransaction.Result execute(SnapshotAggrMaxTx tx) {
        try (Transaction t = db.beginTx()) {
            List<Pair<String, Float>> res = new ArrayList<>();
            HashMap<String, Float> tmp = new HashMap<>();
            int st = tx.getT0();
            int en = tx.getT1();
            try(ResourceIterator<Node> it = db.findNodes(Label.label("TPNODE"))) {
                while(it.hasNext()) {
                    Node tp = it.next();
                    int startTime = (int) tp.getProperty("st_time");
                    int endTime = (int) tp.getProperty("en_time");
                    if (startTime <= en && st <= endTime) {
                        String name = null;
                        for(Relationship r : tp.getRelationships()){
                            Node main = r.getOtherNode(tp);
                            name = (String) main.getProperty("name");
                        }
                        assert name!=null;
                        float v = (Float) tp.getProperty(tx.getP());
                        tmp.merge(name, v, Float::max);
                    }
                }
            }
            t.success();
            tmp.forEach((k, v) -> res.add(Pair.of(k, v)));
            SnapshotAggrMaxTx.Result result = new SnapshotAggrMaxTx.Result();
            result.setPropMaxValue(res);
            return result;
        }
    }

    private AbstractTransaction.Result execute(SnapshotAggrDurationTx tx) {
        try (Transaction t = db.beginTx()) {
            List<Triple<String, Integer, Integer>> res = new ArrayList<>();
            // id, property_value, duration
            HashMap<String, HashMap<Integer, Integer>> buffer = new HashMap<>();
            int t0 = tx.getT0();
            int t1 = tx.getT1();
            try(ResourceIterator<Node> it = db.findNodes(Label.label("TPNODE"))) {
                while(it.hasNext()) {
                    Node tp = it.next();
                    int st = (int) tp.getProperty("st_time");
                    int en = (int) tp.getProperty("en_time");
                    if (t1 >= st && t0 <= en) {
                        int duration;
                        if (t0 >= st && t1 <= en) { // inner
                            duration = t1 - t0;
                        } else if (t0 >= st) { // left
                            duration = en - t0 + 1;
                        } else if (t1 >= en) { // middle
                            if (t1 == en) duration = t1 - st;
                            else duration = en - st + 1;
                        } else { // right
                            duration = t1 - st;
                        }
                        String name = null;
                        for(Relationship r : tp.getRelationships()){
                            Node main = r.getOtherNode(tp);
                            name = (String) main.getProperty("name");
                        }
                        assert name!=null;
                        float value = (float) tp.getProperty(tx.getP());
                        Integer grp = tx.getIntervalStarts().floor((int) value);
                        if(grp==null) continue;
                        HashMap<Integer, Integer> tmp = buffer.get(name);
                        if (tmp != null) {
                            tmp.merge(grp, duration, Integer::sum);
                        } else {
                            tmp = new HashMap<>();
                            tmp.put(grp, duration);
                            buffer.put(name, tmp);
                        }
                    }
                }
            }
            t.success();
            buffer.forEach((key, value) -> value.forEach((k, v) -> res.add(Triple.of(key, k, v))));
            SnapshotAggrDurationTx.Result result = new SnapshotAggrDurationTx.Result();
            result.setStatusDuration(res);
            return result;
        }
    }

    private AbstractTransaction.Result execute(EntityTemporalConditionTx tx) {
        int st = tx.getT0();
        int en = tx.getT1();
        float vMin = tx.getVMin();
        float vMax = tx.getVMax();
        HashSet<String> res = new HashSet<>();
        try (Transaction t = db.beginTx();
             ResourceIterator<Node> it = db.findNodes(Label.label("TPNODE"))) {
            while(it.hasNext()) {
                Node tp = it.next();
                int startTime = (int) tp.getProperty("st_time");
                int endTime = (int) tp.getProperty("en_time");
                float v = (Float) tp.getProperty(tx.getP());
                if (startTime <= en && st <= endTime && vMin<=v && v<=vMax) {
                    String name = null;
                    for(Relationship r : tp.getRelationships()){
                        Node main = r.getOtherNode(tp);
                        name = (String) main.getProperty("name");
                    }
                    assert name!=null;
                    res.add(name);
                }
            }
            t.success();
        }
        EntityTemporalConditionTx.Result result = new EntityTemporalConditionTx.Result();
        result.setEntities(new ArrayList<>(res));
        return result;
    }

}
